from helper import decode_json
from helper import *

import multiprocessing as mp
import threading
from queue import Queue as ThreadQueue
import time
import cv2
import os
import sys
import apriltag
import socket

sys.path.append("../")

from new_arm_control.angle import ArmAngle
from new_arm_control.move import ArmMove
from new_arm_control.func import execute_operation

CUBE_HEIGHT = 3.6
MAX_NUMBER_OF_OBEJECTS = 6

PICK_AND_RELEASE_OFFSET = 6
ARM_Z_OFFSET = 5

WINDOW_W = 640
WINDOW_H = 480
WINDOW_NAME = "Project P: Voice Control Robot Pick and Release"

SERVER_ADDRESS = "localhost"
SERVER_PORT = 12345
server_socket = None

CONFIG_FILE = os.path.expanduser("~/.bc_config/cam_calibration.config")
cfg = load_config(CONFIG_FILE)

arm_base_coordinate_x = cfg["arm_base_coordinate"]["x"]
arm_base_coordinate_y = cfg["arm_base_coordinate"]["y"]
pixel_per_cm_ratio = cfg["pixel_per_cm_ratio"]

# 放置区域范围(7，8,9在右边；10,11,12在左边)（x左正右负，y前正后负）单位cm）
place_position_in_pixels = {
    7: (arm_base_coordinate_x - int(18 * pixel_per_cm_ratio), arm_base_coordinate_y + int(4.5 * pixel_per_cm_ratio)),
    8: (arm_base_coordinate_x - int(18 * pixel_per_cm_ratio), arm_base_coordinate_y + int(17 * pixel_per_cm_ratio)),
    9: (arm_base_coordinate_x + int(18 * pixel_per_cm_ratio), arm_base_coordinate_y + int(4.5 * pixel_per_cm_ratio)),
    10: (arm_base_coordinate_x + int(18 * pixel_per_cm_ratio), arm_base_coordinate_y + int(17 * pixel_per_cm_ratio)),
}


# 计算机械臂在摄像头像素坐标下距离原点的实际距离(x,y,z)
arm_base_pos = [
    *compute_pixel_to_dist(
        arm_base_coordinate_x, arm_base_coordinate_y, pixel_per_cm_ratio
    ),
    0,
]


def process_wrapper(queue, signal):
    def decorator(func):
        def wrapper():
            arm_move = ArmMove()
            arm_angle = ArmAngle(base_pos=arm_base_pos)
            while True:
                if not queue.empty():
                    targets = queue.get()
                    func(arm_move, arm_angle, *targets)
                    signal.set()
                time.sleep(0.1)  # Prevents busy waiting
        return wrapper
    return decorator


# Create a multiprocessing queue
queue = mp.Queue()
signal = mp.Event()
signal.set()


# Create a decorated version of pick_and_release
@process_wrapper(queue, signal)
def wrapped_execute_operation(*args):
    execute_operation(*args)


def socket_handler(sock, command_queue):
    # 处理socket接收到的指令，将其放入命令队列
    while True:
        data, addr = sock.recvfrom(1024)
        command = data.decode()
        command_queue.put(command)

def udp_server(host="localhost", port=12345):
    # 创建UDP服务器socket
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server_socket.bind((host, port))
    return server_socket

def run(min_dist_threshold=10):
    # 主循环，处理摄像头画面、AprilTag检测、命令接收与机械臂控制
    
    process = mp.Process(target=wrapped_execute_operation)
    process.start()
    time.sleep(3)
    
    cap = cv2.VideoCapture(0)
    detector = apriltag.Detector()

    cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
    cv2.resizeWindow(WINDOW_NAME, WINDOW_W, WINDOW_H)

    place_position_in_dist = {
        name: (*compute_pixel_to_dist(x, y, pixel_per_cm_ratio), CUBE_HEIGHT)
        for name, (x, y) in place_position_in_pixels.items()
    }
    
    print(place_position_in_dist)

    target_box_size_in_pixel = int(CUBE_HEIGHT * pixel_per_cm_ratio)
    
    server_socket = udp_server(host=SERVER_ADDRESS, port=SERVER_PORT)
    
    command_queue = ThreadQueue()
    
    threading.Thread(
        target=socket_handler, args=(server_socket, command_queue), daemon=True
    ).start()

    i = 0
    cached_detected_pos = {}
    cached_goal_pos = -1
    pos_dic = {}

    try:
        ### 整个while 是检测函数的主体, 在主体里面插入 pick_and_release 函数进程的控制信号量
        while True:
            ret, frame = cap.read()
            gray_img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

            # 对输入的图像检测 Apriltag
            tag_pos_pairs = []
            results = detector.detect(gray_img)
            for result in results:
                tag_id = result.tag_id
                top_left, top_right, bottom_right, bottom_left = result.corners
                center_x = int((top_left[0] + bottom_right[0]) / 2)
                center_y = int((top_left[1] + bottom_right[1]) / 2)
                tag_inform = (tag_id, center_x, center_y)
                tag_pos_pairs.append(tag_inform)

            # 可视化检测结果
            for tag in tag_pos_pairs:
                draw_tag(frame, tag[0], tag[1:])
                cached_detected_pos[tag[0]] = (*compute_pixel_to_dist(tag[1], tag[2], pixel_per_cm_ratio), CUBE_HEIGHT)

            # 可视化放置位置
            frame = draw_target_pos(frame, place_position_in_pixels, width=target_box_size_in_pixel)
            cv2.putText(
                frame,
                f"Jobs Count: {i}",
                (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX,
                1,
                (0, 255, 0),
                2,
            )

            # 当前正在抓取中,不执行任何其它抓取判断
            if not signal.is_set():
                time.sleep(0.03)
                if not command_queue.empty():
                    command = command_queue.get()
                    print(f"\n接收到的command:\n {command}")
                    print(f"机械臂正在任务中, 忽略此次接收到的指令, 请重新发送")
            else:
                # 从信息队列中读取控制指令
                if not command_queue.empty():
                    command = command_queue.get()
                    print(f"\n接收到的command:\n {command}")
                    command = decode_json(command)
                    # 处理长指令
                    pick_pos = command.get("pick_target", -1)
                    place_pos = command.get("place_target", -1)
                    if pick_pos != -1 and place_pos != -1:
                        # 验证位置范围
                        if 1 <= pick_pos <= 4 and 7 <= place_pos <= 10:
                            if pick_pos in cached_detected_pos and place_pos in place_position_in_dist:
                                pick_x, pick_y, pick_z = cached_detected_pos[pick_pos]
                                place_x, place_y, place_z = place_position_in_dist[place_pos]
                                print(
                                    f"将坐标发送到队列: \n\t抓取目标为 {pick_pos}\n\t放置目标为 {place_pos}"
                                )
                                queue.put(
                                    (
                                        pick_x, pick_y, pick_z,
                                        "pick",
                                        PICK_AND_RELEASE_OFFSET
                                    )
                                )
                                queue.put(
                                    (
                                        place_x, place_y, place_z,
                                        "release",
                                        PICK_AND_RELEASE_OFFSET
                                    )
                                )
                                i += 1
                                signal.clear()
                            else:
                                print(f"没有找到任务目标: 抓取 {pick_pos}, 放置 {place_pos}")
                        else:
                            print(f"位置超出预设范围: 抓取 {pick_pos}, 放置 {place_pos}")
                    else:
                        print("指令格式错误，缺少抓取或放置目标")

            cv2.imshow(WINDOW_NAME, frame)
            if cv2.waitKey(10) & 0xFF == ord('q'):
                break
    except KeyboardInterrupt:
        print("Process interrupted by user")
    finally:
        process.terminate()
        process.join()


if __name__ == "__main__":
    min_dist_threshold = 15
    run(min_dist_threshold)